{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# colabdock配置文件与约束设置\n",
    "\n",
    "## 相关链接\n",
    "\n",
    "- [论文预印本:S. Feng et al., “ColabDock: inverting AlphaFold structure prediction model for protein-protein docking with experimental restraints,” Jul. 04, 2023, bioRxiv. doi: 10.1101/2023.07.04.547599.](https://www.biorxiv.org/content/10.1101/2023.07.04.547599v1)\n",
    "- [论文:S. Feng et al., “Integrated structure prediction of protein–protein docking with experimental restraints using ColabDock,” Nat Mach Intell, vol. 6, no. 8, pp. 924–935, Aug. 2024, doi: 10.1038/s42256-024-00873-z.](https://www.nature.com/articles/s42256-024-00873-z)\n",
    "- [ColabDock git仓库](https://github.com/JeffSHF/ColabDock.git)\n",
    "- 我修改的git仓库: [https://gitee.com/regentsai/colab-dock-fork.git](https://gitee.com/regentsai/colab-dock-fork.git)\n",
    "\n",
    "\n",
    "## 配置文件示例"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 442,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    # 保存预测结果的目录\n",
    "    'save_path': 'data/results',\n",
    "\n",
    "    ###########################################################################################################\n",
    "    # 单体模板和复合物参考结构\n",
    "    ###########################################################################################################\n",
    "    # 单体模板的pdb结构\n",
    "    'template': 'data/4HFF.pdb',\n",
    "\n",
    "    # [可选]复合物参考结构,用于计算RMSD,如果没有则设置为None.\n",
    "    'native': './data/4HFF.pdb',\n",
    "\n",
    "    # 对接的链号\n",
    "    # 顺序决定对接时拼接的顺序.\n",
    "    'chains': 'A,B',\n",
    "\n",
    "    # 固定的链号,让指定链之间的相对位置与模板保持不变,没有要求则设置为None\n",
    "    # 例如:\n",
    "    #     'fixed_chains': ['A,B', 'C,D']\n",
    "    #     A-B链的相对位置不变,C-D链的相对位置不变.\n",
    "    #     'fixed_chains': ['A,B'],\n",
    "    'fixed_chains': None,\n",
    "\n",
    "    ###########################################################################################################\n",
    "    # 实验约束\n",
    "    ###########################################################################################################\n",
    "    # 实验约束的距离阈值,默认为8Å.\n",
    "    # 由于AF2的定义限制,阈值需要在2-22Å之间\n",
    "    'res_thres': 8.0,\n",
    "\n",
    "    # 1v1 约束\n",
    "    # 描述:\n",
    "    #     距离在给定阈值以下的2个异链残基,没有则为None.\n",
    "    #     多对1v1约束用[]进行罗列.\n",
    "    #     残基序号指的是它在完整序列中的序号,序号从1开始.\n",
    "    #         完整序列由所有指定的链的序号拼接,拼接顺序由`chains`定义.\n",
    "    # 示例:\n",
    "    #     'rest_1v1': [55,90]:整条链上的第55号残基和90号残基距离低于规定阈值\n",
    "    #     'rest_1v1': [[78,198],[20,50]]:整条链上的第78号残基与198号残基距离/第20号残基与50号残基距离均低于规定阈值\n",
    "    'rest_1v1': [79, 199],\n",
    "\n",
    "    # 1vN 约束\n",
    "    # 描述:\n",
    "    #     距离在给定阈值以下的1个残基与另1组残基,没有则为None\n",
    "    # 示例:\n",
    "    #     'rest_1vN': [36,list(range(160,171))+[178,190]]\n",
    "    #     整条链上36号残基与[160~170,178,190]号残基中至少有1个距离满足约束.注意range是左闭右开的\n",
    "    'rest_1vN': None,\n",
    "\n",
    "    # MvN 约束\n",
    "    # 描述:\n",
    "    #     包含若干个1vN约束,只有给定数量的残基对满足约束,没有则为None\n",
    "    #     多个MvN约束用[]进行罗列\n",
    "    # 示例:\n",
    "    #     'rest_MvN': [[10, list(range(160, 170))],\n",
    "    #                  [78, list(range(160, 170))],\n",
    "    #                  [120, list(range(160, 170))],\n",
    "    #                  2]\n",
    "    #     3个给定的1vN约束,其中只有2个满足约束.\n",
    "    # 也支持从pkl数据文件中导入(由extract_rest.py生成)\n",
    "    # 'rest_MvN': joblib.load('./protein/1AHW/rest_MvN.pkl'),\n",
    "    'rest_MvN': None,\n",
    "\n",
    "    # 排斥距离阈值,两个残基C-beta原子的距离需要超过该距离,需要在2-22埃范围内\n",
    "    'rep_thres': 8.0,\n",
    "\n",
    "    # 排斥约束,与1v1约束规定类似,但距离要超过阈值\n",
    "    # 示例:\n",
    "    #     'rest_rep: [154, 250]\n",
    "    #     整条链上第154,250号残基距离超过阈值\n",
    "    'rest_rep': None,\n",
    "\n",
    "    ###########################################################################################################\n",
    "    # 优化参数\n",
    "    ###########################################################################################################\n",
    "    # 基于片段的优化切割长度,若没有则为None.\n",
    "    # 对于长链序列,建议值为200\n",
    "    # 可以减少显存消耗,但可能陷入局部最优化\n",
    "    'crop_len': 200,\n",
    "\n",
    "    # 重复预测的轮次数(rounds)\n",
    "    # 增大数值会提高精度,但会增加模拟时间\n",
    "    # 论文的案例重复5次\n",
    "    'rounds': 1,\n",
    "\n",
    "    # 每轮次反向传播的次数(steps),让结构收敛\n",
    "    # 基于片段的优化建议使用更大的数值,如150(论文案例为100)\n",
    "    # 否则建议设置为50\n",
    "    'steps': 100,\n",
    "\n",
    "    # 每n次反向传播保存1次构象.\n",
    "    # 对基于片段的优化很有用,以缩减生成结构的文件大小,建议设置为3以上\n",
    "    # 普通优化过程设置为1即可\n",
    "    'save_every_n_step': 10,\n",
    "\n",
    "    ###########################################################################################################\n",
    "    # AF2 模型\n",
    "    ###########################################################################################################\n",
    "    # AF2 权重参数所在路径\n",
    "    'data_dir': '/home/regen/projects/af_para',\n",
    "    # 使用AF2-multimer,设为False则使用AF2\n",
    "    'use_multimer': True,\n",
    "\n",
    "    # 是否使用bfloat浮点数精度(速度上升,精度下降)\n",
    "    'bfloat': True,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "补充要点:\n",
    "\n",
    "1. 源仓库默认在源代码根目录下运行,必须有一个config.py配置文件\n",
    "2. 我的修改版本可以在任意位置下运行main.py,且配置文件的位置和名称可以调整\n",
    "3. 保存结果/模板/参考结构/约束/AF参数选项中使用的路径如果是相对路径,相对路径必须以运行python脚本的当前路径为基准,无论config.py所在的位置\n",
    "4. 对于残基号不连续的蛋白结构模板,个人建议先补齐残基(使用Alphafold,pdbfixer等).虽然生成的结构能够为丢失的残基保留空位,但残基序号是连续的,可能在位置编码中引入噪声\n",
    "\n",
    "![](data/4HFF_missing.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. `config.py和main.py`接受的残基序号由指定的链号顺序和他们的长度有关,与原始编号无关.\n",
    "   1. 例如指定A,B,C链对接,A/B/C长度分别为10,20,40,则在总序列中各链的序号为1-10,11-30,31-70.\n",
    "   2. 模板结构中各链中残基序号不连续,不会影响在AF2框架中的总序列序号\n",
    "   3. 若A链的第3个残基与C链的第5个残基满足1v1约束,则1v1约束定义为[3,35]\n",
    "   4. 若对接过程不包含B链,则1v1约束定义为[3,15]\n",
    "   5. `extract_rest.py`自动提取的约束编号满足以上要求,但无法快速了解提取残基的原始编号\n",
    "   6. 用户从实验或文献中获取的约束信息,需要按以上算法进行计算,才能手动输入约束,对用户不友好,建议改善"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "ColabDockRoot=os.path.expanduser('~/git_develop/ColabDock') #克隆路径自行调整"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 443,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/regen/ana/envs/colabdock/lib/python3.8/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n",
      "restraints:\n",
      "\t1v1 restraints:\n",
      "\t\t[79, 199]\n",
      "\tno 1vN restraints provided.\n",
      "\tno MvN restraints provided.\n",
      "\tno repulsive restraints provided.\n",
      "\n",
      "Optimization losses include:\n",
      "\t1v1 restraint loss, distogram loss, pLDDT, and ipAE.\n",
      "Colabdock will work in segment based mode.\n",
      "\n",
      "Start optimization\n",
      "/home/regen/git_develop/ColabDock/colabdesign/af/alphafold/model/geometry/struct_of_array.py:136: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.\n",
      "  flat_array_like, inner_treedef = jax.tree_flatten(array_like)\n",
      "/home/regen/git_develop/ColabDock/colabdesign/af/alphafold/model/geometry/struct_of_array.py:209: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead.\n",
      "  value_dict[array_field] = jax.tree_unflatten(\n",
      "/home/regen/git_develop/ColabDock/colabdesign/af/alphafold/model/mapping.py:50: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.\n",
      "  values_tree_def = jax.tree_flatten(values)[1]\n",
      "/home/regen/git_develop/ColabDock/colabdesign/af/alphafold/model/mapping.py:54: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead.\n",
      "  return jax.tree_unflatten(values_tree_def, flat_axes)\n",
      "/home/regen/git_develop/ColabDock/colabdesign/af/alphafold/model/mapping.py:128: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.\n",
      "  flat_sizes = jax.tree_flatten(in_sizes)[0]\n",
      "1 models [1] recycles 0 hard 0 soft 0 temp 1 loss 2.530 plddt 0.279 i_pae 0.603 dgram_cce 1.628\n",
      "2 models [1] recycles 0 hard 0 soft 0 temp 1 loss 14.226 plddt 0.193 i_pae 0.564 dgram_cce 1.424 rest_1v1 6.007\n",
      "3 models [0] recycles 0 hard 0 soft 0 temp 1 loss 2.267 plddt 0.227 i_pae 0.545 dgram_cce 1.460\n",
      "4 models [1] recycles 0 hard 0 soft 0 temp 1 loss 2.528 plddt 0.236 i_pae 0.635 dgram_cce 1.627\n",
      "5 models [1] recycles 0 hard 0 soft 0 temp 1 loss 2.152 plddt 0.140 i_pae 0.467 dgram_cce 1.394\n",
      "6 models [0] recycles 0 hard 0 soft 0 temp 1 loss 11.300 plddt 0.188 i_pae 0.662 dgram_cce 1.382 rest_1v1 4.571\n",
      "7 models [1] recycles 0 hard 0 soft 0 temp 1 loss 2.216 plddt 0.074 i_pae 0.191 dgram_cce 1.280 rest_1v1 0.135\n",
      "8 models [0] recycles 0 hard 0 soft 0 temp 1 loss 2.669 plddt 0.123 i_pae 0.272 dgram_cce 1.372 rest_1v1 0.286\n",
      "9 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.914 plddt 0.047 i_pae 0.139 dgram_cce 1.250 rest_1v1 0.010\n",
      "10 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.971 plddt 0.056 i_pae 0.141 dgram_cce 1.262 rest_1v1 0.029\n",
      "11 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.905 plddt 0.064 i_pae 0.229 dgram_cce 1.250\n",
      "12 models [1] recycles 0 hard 0 soft 0 temp 1 loss 2.098 plddt 0.153 i_pae 0.634 dgram_cce 1.346\n",
      "13 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.572 plddt 0.022 i_pae 0.103 dgram_cce 1.033 rest_1v1 0.005\n",
      "14 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.715 plddt 0.046 i_pae 0.128 dgram_cce 1.132\n",
      "15 models [0] recycles 0 hard 0 soft 0 temp 1 loss 2.083 plddt 0.191 i_pae 0.602 dgram_cce 1.336\n",
      "16 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.563 plddt 0.024 i_pae 0.098 dgram_cce 1.030 rest_1v1 0.003\n",
      "17 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.439 plddt 0.017 i_pae 0.087 dgram_cce 0.952 rest_1v1 0.001\n",
      "18 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.671 plddt 0.033 i_pae 0.093 dgram_cce 1.106\n",
      "19 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.402 plddt 0.019 i_pae 0.089 dgram_cce 0.926 rest_1v1 0.001\n",
      "20 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.592 plddt 0.029 i_pae 0.089 dgram_cce 1.054\n",
      "21 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.675 plddt 0.046 i_pae 0.161 dgram_cce 1.103\n",
      "22 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.606 plddt 0.039 i_pae 0.110 dgram_cce 1.061\n",
      "23 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.469 plddt 0.018 i_pae 0.100 dgram_cce 0.964 rest_1v1 0.005\n",
      "24 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.378 plddt 0.017 i_pae 0.087 dgram_cce 0.911 rest_1v1 0.001\n",
      "25 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.568 plddt 0.041 i_pae 0.146 dgram_cce 1.033\n",
      "26 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.517 plddt 0.036 i_pae 0.121 dgram_cce 1.001\n",
      "27 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.334 plddt 0.012 i_pae 0.084 dgram_cce 0.882 rest_1v1 0.001\n",
      "28 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.274 plddt 0.012 i_pae 0.080 dgram_cce 0.843 rest_1v1 0.000\n",
      "29 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.477 plddt 0.025 i_pae 0.079 dgram_cce 0.977\n",
      "30 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.253 plddt 0.014 i_pae 0.081 dgram_cce 0.828 rest_1v1 0.001\n",
      "31 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.248 plddt 0.011 i_pae 0.078 dgram_cce 0.825 rest_1v1 0.001\n",
      "32 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.180 plddt 0.011 i_pae 0.076 dgram_cce 0.780 rest_1v1 0.001\n",
      "33 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.421 plddt 0.019 i_pae 0.083 dgram_cce 0.941\n",
      "34 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.308 plddt 0.021 i_pae 0.088 dgram_cce 0.865\n",
      "35 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.149 plddt 0.009 i_pae 0.073 dgram_cce 0.760 rest_1v1 0.001\n",
      "36 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.265 plddt 0.016 i_pae 0.080 dgram_cce 0.836 rest_1v1 0.001\n",
      "37 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.211 plddt 0.013 i_pae 0.077 dgram_cce 0.801 rest_1v1 0.001\n",
      "38 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.547 plddt 0.149 i_pae 0.508 dgram_cce 0.988\n",
      "39 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.169 plddt 0.012 i_pae 0.075 dgram_cce 0.772 rest_1v1 0.001\n",
      "40 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.181 plddt 0.013 i_pae 0.075 dgram_cce 0.781 rest_1v1 0.000\n",
      "41 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.136 plddt 0.013 i_pae 0.075 dgram_cce 0.751 rest_1v1 0.001\n",
      "42 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.318 plddt 0.035 i_pae 0.172 dgram_cce 0.865\n",
      "43 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.465 plddt 0.040 i_pae 0.101 dgram_cce 0.967\n",
      "44 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.165 plddt 0.013 i_pae 0.073 dgram_cce 0.770 rest_1v1 0.001\n",
      "45 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.318 plddt 0.030 i_pae 0.119 dgram_cce 0.869\n",
      "46 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.266 plddt 0.014 i_pae 0.069 dgram_cce 0.838\n",
      "47 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.131 plddt 0.011 i_pae 0.071 dgram_cce 0.748 rest_1v1 0.000\n",
      "48 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.084 plddt 0.011 i_pae 0.071 dgram_cce 0.717 rest_1v1 0.000\n",
      "49 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.234 plddt 0.019 i_pae 0.084 dgram_cce 0.816\n",
      "50 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.105 plddt 0.013 i_pae 0.073 dgram_cce 0.730 rest_1v1 0.001\n",
      "51 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.098 plddt 0.010 i_pae 0.069 dgram_cce 0.726 rest_1v1 0.000\n",
      "52 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.062 plddt 0.011 i_pae 0.070 dgram_cce 0.702 rest_1v1 0.001\n",
      "53 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.057 plddt 0.010 i_pae 0.067 dgram_cce 0.699 rest_1v1 0.000\n",
      "54 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.067 plddt 0.011 i_pae 0.067 dgram_cce 0.706 rest_1v1 0.000\n",
      "55 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.016 plddt 0.009 i_pae 0.065 dgram_cce 0.672 rest_1v1 0.000\n",
      "56 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.366 plddt 0.037 i_pae 0.207 dgram_cce 0.895\n",
      "57 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.369 plddt 0.019 i_pae 0.083 dgram_cce 0.906\n",
      "58 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.224 plddt 0.027 i_pae 0.087 dgram_cce 0.768 rest_1v1 0.030\n",
      "59 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.292 plddt 0.019 i_pae 0.074 dgram_cce 0.855\n",
      "60 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.257 plddt 0.014 i_pae 0.077 dgram_cce 0.832\n",
      "61 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.234 plddt 0.019 i_pae 0.079 dgram_cce 0.816\n",
      "62 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.075 plddt 0.010 i_pae 0.067 dgram_cce 0.711 rest_1v1 0.000\n",
      "63 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.179 plddt 0.012 i_pae 0.063 dgram_cce 0.781\n",
      "64 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.163 plddt 0.010 i_pae 0.069 dgram_cce 0.768 rest_1v1 0.001\n",
      "65 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.277 plddt 0.022 i_pae 0.085 dgram_cce 0.844\n",
      "66 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.199 plddt 0.020 i_pae 0.107 dgram_cce 0.791\n",
      "67 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.218 plddt 0.014 i_pae 0.066 dgram_cce 0.807\n",
      "68 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.235 plddt 0.023 i_pae 0.093 dgram_cce 0.815\n",
      "69 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.157 plddt 0.015 i_pae 0.074 dgram_cce 0.765\n",
      "70 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.148 plddt 0.009 i_pae 0.063 dgram_cce 0.761\n",
      "71 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.098 plddt 0.007 i_pae 0.065 dgram_cce 0.726 rest_1v1 0.000\n",
      "72 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.075 plddt 0.011 i_pae 0.066 dgram_cce 0.711 rest_1v1 0.000\n",
      "73 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.039 plddt 0.008 i_pae 0.065 dgram_cce 0.687 rest_1v1 0.001\n",
      "74 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.075 plddt 0.012 i_pae 0.068 dgram_cce 0.711 rest_1v1 0.000\n",
      "75 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.009 plddt 0.007 i_pae 0.064 dgram_cce 0.667 rest_1v1 0.001\n",
      "76 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.034 plddt 0.011 i_pae 0.066 dgram_cce 0.684 rest_1v1 0.000\n",
      "77 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.159 plddt 0.017 i_pae 0.094 dgram_cce 0.766\n",
      "78 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.173 plddt 0.017 i_pae 0.080 dgram_cce 0.776\n",
      "79 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.048 plddt 0.007 i_pae 0.064 dgram_cce 0.693 rest_1v1 0.000\n",
      "80 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.056 plddt 0.010 i_pae 0.065 dgram_cce 0.699 rest_1v1 0.000\n",
      "81 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.146 plddt 0.018 i_pae 0.085 dgram_cce 0.757\n",
      "82 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.026 plddt 0.008 i_pae 0.064 dgram_cce 0.678 rest_1v1 0.001\n",
      "83 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.119 plddt 0.013 i_pae 0.073 dgram_cce 0.740 rest_1v1 0.000\n",
      "84 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.006 plddt 0.007 i_pae 0.064 dgram_cce 0.665 rest_1v1 0.000\n",
      "85 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.306 plddt 0.099 i_pae 0.318 dgram_cce 0.843\n",
      "86 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.041 plddt 0.012 i_pae 0.066 dgram_cce 0.688 rest_1v1 0.001\n",
      "87 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.120 plddt 0.010 i_pae 0.063 dgram_cce 0.742\n",
      "88 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.048 plddt 0.008 i_pae 0.060 dgram_cce 0.694\n",
      "89 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.051 plddt 0.011 i_pae 0.066 dgram_cce 0.695 rest_1v1 0.001\n",
      "90 models [1] recycles 0 hard 0 soft 0 temp 1 loss 0.992 plddt 0.007 i_pae 0.060 dgram_cce 0.656 rest_1v1 0.001\n",
      "91 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.142 plddt 0.016 i_pae 0.077 dgram_cce 0.755\n",
      "92 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.037 plddt 0.010 i_pae 0.064 dgram_cce 0.686 rest_1v1 0.000\n",
      "93 models [0] recycles 0 hard 0 soft 0 temp 1 loss 0.969 plddt 0.010 i_pae 0.061 dgram_cce 0.641 rest_1v1 0.000\n",
      "94 models [1] recycles 0 hard 0 soft 0 temp 1 loss 0.968 plddt 0.007 i_pae 0.059 dgram_cce 0.641 rest_1v1 0.000\n",
      "95 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.082 plddt 0.009 i_pae 0.064 dgram_cce 0.717\n",
      "96 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.022 plddt 0.008 i_pae 0.063 dgram_cce 0.677\n",
      "97 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.030 plddt 0.007 i_pae 0.056 dgram_cce 0.683\n",
      "98 models [1] recycles 0 hard 0 soft 0 temp 1 loss 1.005 plddt 0.006 i_pae 0.057 dgram_cce 0.665 rest_1v1 0.000\n",
      "99 models [0] recycles 0 hard 0 soft 0 temp 1 loss 1.256 plddt 0.026 i_pae 0.094 dgram_cce 0.829\n",
      "100 models [1] recycles 0 hard 0 soft 0 temp 1 loss 0.958 plddt 0.006 i_pae 0.057 dgram_cce 0.634 rest_1v1 0.000\n",
      "generation stage inference:\n",
      "infer epoch 1\n",
      "100%|███████████████████████████████████████████| 10/10 [00:26<00:00,  2.70s/it]\n",
      "prediction stage inference:\n",
      "infer epoch 1\n",
      "100%|███████████████████████████████████████████| 10/10 [00:20<00:00,  2.03s/it]\n",
      "1st_best structure:\n",
      "\trmsd: 1.103, iptm: 0.940, 1 out of 1 restraints are satisfied.\n",
      "2nd_best structure:\n",
      "\trmsd: 1.133, iptm: 0.938, 1 out of 1 restraints are satisfied.\n",
      "3rd_best structure:\n",
      "\trmsd: 1.155, iptm: 0.934, 1 out of 1 restraints are satisfied.\n",
      "4th_best structure:\n",
      "\trmsd: 16.453, iptm: 0.101, 0 out of 1 restraints are satisfied.\n",
      "5th_best structure:\n",
      "\trmsd: 17.259, iptm: 0.124, 0 out of 1 restraints are satisfied.\n"
     ]
    }
   ],
   "source": [
    "!python $ColabDockRoot/main.py -c data/config_4HFF.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 约束提取\n",
    "\n",
    "源代码仓库下的extract_rest.py进行随机约束提取."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/regen/ana/envs/colabdock/lib/python3.8/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n",
      "usage: extract_rest.py [-h] [--rest_thres REST_THRES] [--N N] [--M M]\n",
      "                       [--num NUM] [--save_path SAVE_PATH] [--init INIT]\n",
      "                       [--verbose]\n",
      "                       pdb chains chains_rest rest_type\n",
      "\n",
      "randomly extract restraints from a given complex structure. An example:\n",
      "`python extract_rest.py ./protein/4INS4/PDB/native.pdb A,B,C,D A,D 1v1 -i\n",
      "45,12`\n",
      "\n",
      "positional arguments:\n",
      "  pdb                   input complex structure\n",
      "  chains                docking chains, this parameter should be the same as\n",
      "                        that in the config.py file\n",
      "  chains_rest           sample restraints between which two chains\n",
      "  rest_type             sampled restraints type. Please choose one type from\n",
      "                        1v1, 1vN, MvN, and repulsive\n",
      "\n",
      "optional arguments:\n",
      "  -h, --help            show this help message and exit\n",
      "  --rest_thres REST_THRES\n",
      "                        distance threshold of the restraints. Default 8.0\n",
      "  --N N                 In the sampled 1vN restraint, the distance between a\n",
      "                        residue and at least one of N residues is below a\n",
      "                        certain value. Default 5.\n",
      "  --M M                 number of 1vN restraints in the sampled MvN restraint.\n",
      "                        Default 2.\n",
      "  --num NUM             number of sampled restraints. Default 1.\n",
      "  --save_path SAVE_PATH\n",
      "                        the file to save the sampled restraints using joblib\n",
      "                        package. Default None prints the restraints to stdout.\n",
      "  --init INIT, -i INIT  initial residue indices in template pdb file.Accept\n",
      "                        1v1/repulsive restraint(s) only and ignore other\n",
      "                        restraint arguments if it is set.\n",
      "  --verbose, -v         print more information about the restrants.Double\n",
      "                        check the initial residue names since pdb segment ids\n",
      "                        can be different from their papers.\n"
     ]
    }
   ],
   "source": [
    "!python $ColabDockRoot/extract_rest.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "用法: extract_rest.py [-h] [--rest_thres REST_THRES] [--N N] [--M M]\n",
    "                       [--num NUM] [--save_path SAVE_PATH]\n",
    "                       pdb chains chains_rest rest_type\n",
    "\n",
    "- 位置参数:\n",
    "  - pdb:输入的复合物结构\n",
    "  - chains:要对接的链,链号和顺序应与config.py一致\n",
    "  - chains_rest:约束所在的链号,目前只能指定2个单独的链,不接受亚单位采样\n",
    "  - rest_type:采样的约束类型,1v1,1vN,MvN或repulsive\n",
    "- 可选参数:\n",
    "  - --rest_thres REST_THRES:距离阈值,默认为8\n",
    "  - --N N:1vN约束中的N,默认为5\n",
    "  - --M M:MvN约束中的M,默认为2\n",
    "  - --num NUM:采样约束的数量,默认为1\n",
    "  - --save_path SAVE_PATH:保存约束文件到joblib二进制文件中,拓展名应为pkl,默认不保存,仅将约束打印到标准输出.\n",
    "  - --init/-i INIT_IDS:1v1(排斥)约束的残基对在模板中的编号,默认从模板结构中自动提取,如果指定模板编号则不考虑可选的残基约束参数.接受的输入形式包括:\n",
    "    - 45,12\n",
    "    - [45,12]\n",
    "    - [[45,12]]\n",
    "    - [45,12],[47,22],[47,20],...\n",
    "    - [[45,12],[47,22],[47,20],...]\n",
    "    - 出于安全考虑,仅支持python字面量,不支持range,zip等函数调用\n",
    "  - --verbose/-v:额外打印约束残基在模板中的编号和残基名,方便交叉检查残基信息.默认不打印额外信息.\n",
    "- 示例命令:\n",
    "\n",
    "```py\n",
    "python extract_rest.py ./protein/4INS4/PDB/native.pdb A,B,C,D A,D 1v1 # 自动提取\n",
    "python extract_rest.py ./protein/4INS4/PDB/native.pdb A,B,C,D A,D 1v1 -i 45,12 -v # 手动指定,编号仅作演示\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 448,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/regen/ana/envs/colabdock/lib/python3.8/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n",
      "The sampled MvN restraints:\n",
      "[[[84, [323, 315, 338, 340, 395]], [107, [362, 248, 270, 236, 306]], 1], [[73, [293, 265, 266, 258, 236]], [43, [285, 356, 359, 243, 293]], 1], [[43, [289, 414, 229, 421, 228]], [5, [384, 252, 417, 361, 432]], 1]]\n",
      "\n",
      "The initial restraints:\n",
      "                                                                         0  \\\n",
      "0  [(84, ALA), [(96, GLU), (88, GLU), (111, GLN), (113, PHE), (168, ASN)]]   \n",
      "1    [(107, ASN), [(135, ASN), (21, GLU), (43, LYS), (9, LEU), (79, ASN)]]   \n",
      "\n",
      "                                                                        1  \\\n",
      "0     [(73, LEU), [(66, GLN), (38, LYS), (39, SER), (31, TYR), (9, LEU)]]   \n",
      "1  [(43, SER), [(58, ASP), (129, THR), (132, ARG), (16, PHE), (66, GLN)]]   \n",
      "\n",
      "                                                                         2  \n",
      "0     [(43, SER), [(62, LYS), (187, GLN), (2, ASN), (194, THR), (1, THR)]]  \n",
      "1  [(5, THR), [(157, SER), (25, LYS), (190, ILE), (134, ASN), (205, GLU)]]  \n"
     ]
    }
   ],
   "source": [
    "!python $ColabDockRoot/extract_rest.py ./data/1AHW.pdb A,B,C A,C MvN --num=3 --save_path=./data/1AHWMvN_3.pkl -v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/regen/ana/envs/colabdock/lib/python3.8/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n",
      "The sampled 1v1 restraints:\n",
      "[[84, 323], [84, 315], [73, 293]]\n",
      "\n",
      "The initial restraints:\n",
      "0    [(84, ALA), (96, GLU)]\n",
      "1    [(84, ALA), (88, GLU)]\n",
      "2    [(73, LEU), (66, GLN)]\n",
      "dtype: object\n"
     ]
    }
   ],
   "source": [
    "!python $ColabDockRoot/extract_rest.py ./data/1AHW.pdb A,B,C A,C 1v1 --save_path=./data/1AHW1v1_3.pkl -i [84,96],[84,88],[73,66] -v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -r data/results_4HFF\n",
    "!rm -r data/1AHW*3.pkl"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "colabdock",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
